#This script is working on PID control,slidebar mass should set to 0 for fix on GND
import pybullet as p  
#import pybullet_data  
import numpy as np  
import time
import numpy as np  
from scipy.optimize import minimize

# 初始化PyBullet  
p.connect(p.GUI)  
p.setGravity(0,0,-9.8)
#p.setAdditionalSearchPath(pybullet_data.getData_path())  
p.setAdditionalSearchPath("D:\\f\\Open Source\\Prj_Robdog\\Simulation_Env\\bullet3\\data")
# 加载环境  
plane_id = p.loadURDF("plane.urdf") 

# 加载 CartPole 模型
# cartpole_start_pos = [0, 0, 0.5]
# cartpole_start_orientation = p.getQuaternionFromEuler([0, 0, 0])
# cartpole_id = p.loadURDF("cartpole.urdf", cartpole_start_pos, cartpole_start_orientation)
# # 设置初始状态
# p.resetBasePositionAndOrientation(cartpole_id, cartpole_start_pos, cartpole_start_orientation)
# p.resetJointState(cartpole_id, 0, 0, 0)  # 设置关节的初始角度和速度

#cartpole_id = p.loadURDF("cartpole.urdf", [0, 0, 0.5]) # use system file 
#this urdf working on MPC
#cartpole_id = p.loadURDF(".\\cartpole.urdf", [0, 0, 0.5])  # use local file, udpate pole mass  
cartpole_id = p.loadURDF(".\\my_cartpole.urdf", [0, 0, 0.5]) 
# Adjust friction for the cart and the plane  
# p.changeDynamics(cartpole_id, 0, lateralFriction=0.1)  # Cart base with lower lateral friction  
# p.changeDynamics(cartpole_id, 1, lateralFriction=0.05)  # pole with lower lateral friction  
# p.changeDynamics(plane_id, 0, lateralFriction=0.05)     # Plane with lower lateral friction 

# Adjust friction for the cart and the plane  
p.changeDynamics(cartpole_id, -1, lateralFriction=0.1)  # Cart base with lower lateral friction 
p.changeDynamics(cartpole_id, 0, lateralFriction=0.1)  # Cart base with lower lateral friction  
p.changeDynamics(cartpole_id, 1, lateralFriction=0.1)  # pole with lower lateral friction 
p.changeDynamics(plane_id, -1, lateralFriction=0.1)     # Plane with lower lateral friction  
p.changeDynamics(plane_id, 0, lateralFriction=0.1)     # Plane with lower lateral friction 
# Set damping for the cart and pole (example values)  
p.changeDynamics(cartpole_id, -1, linearDamping=0.0, angularDamping=0.0)  # For the cart  
p.changeDynamics(cartpole_id, 0, linearDamping=0.0, angularDamping=0.0)  # For the cart  
p.changeDynamics(cartpole_id, 1, linearDamping=0.0, angularDamping=0.0)  # For the pole

# turn off real time simulation, make simulation faster
p.setRealTimeSimulation(0)
#every step is xx s, to make simulation faster, default is 1/240 = 5ms
p.setTimeStep(0.005) # 50 ms

# Alternatively, if the pole is a child of the cart:  
# Set joint state  
initial_angle = -1  # or any appropriate angle  
p.resetJointState(cartpole_id, 1, initial_angle)  

# 获取所有链接的信息
links_info = p.getVisualShapeData(cartpole_id)
for link_info in links_info:
    print("link info:", link_info)
 
#target_angle = 0  # 目标角度（竖直向上）  


#K = 1000  # 控制增益  

def get_state(): 
    #should not use getBasePostion here, base is slidbar, use get joint position
    #pybullet use 1st link as base link 
    #pos, _ = p.getBasePositionAndOrientation(cartpole_id) 
    # 获取 `cart` 的位置和姿态
    cart_link_index = 0  # 假设 `cart` 的链接索引为 0
    cart_pos,_ = p.getLinkState(cartpole_id, cart_link_index)[:2] #get first 2 array member
    
    #print ("pos in get_state", cart_pos) # cart_pos is 3 element,xyz
    #pos = p.getBasePositionAndOrientation(cartpole_id)
    # Get linear velocity  
    cart_velocity = p.getBaseVelocity(cartpole_id)  #linear and angular velocities
    cart_linear_velocity = cart_velocity[0]  # This provides the (vx, vy, vz) linear velocity  

    # Assuming you are only interested in the x-direction for the cart  
    cart_velocity_x = cart_linear_velocity[0]  
    # print("Cart Velocity:", cart_velocity)  
    # print ("pos in get_state", pos[0]) 
    angle = p.getJointState(cartpole_id, 1)[0]  # 杆的角度  
    angular_velocity = p.getJointState(cartpole_id, 1)[1]
    return np.array([cart_pos[0],cart_velocity_x, angle,angular_velocity])  #x, x_dot, theta, theta_dot


def pid_control(state, target_angle, dt):  
    global integral, previous_error ,integral_max, integral_min
    
    x, theta = state  # Unpack the state  
    angle_error = -target_angle +  theta  # Calculate the error  
    
    # Proportional term  
    P_term = Kp * angle_error  
    
    # Integral term  
    integral += angle_error * dt 
    #avoid integral windup 
    if integral > integral_max:  
        integral = integral_max  
    elif integral < integral_min:  
        integral = integral_min 

    
    I_term = Ki * integral  
    
    # Derivative term  
    derivative = (angle_error - previous_error) / dt  
    D_term = Kd * derivative  
    
    # Calculate control input  
    control_input = P_term + I_term + D_term  
    
    # Clip the control input to prevent excessive commands  
    return np.clip(control_input, -2000, 2000)  



# System dynamics (Cart-Pole Model)  
# Example usage  

def dynamics(state, control_input, dt):  
    x, x_dot, theta, theta_dot = state  

    # Simplified dynamics model for a cart-pole  
    m_cart = 1.0  # Mass of the cart  
    m_pole = 10.0  # Mass of the pole  
    g = 9.81  # Acceleration due to gravity  
    length = 0.5  # Length of the pole  
    
    angle = get_state()[2]
    #print("angle in dynamics", angle)
    
   #somehow the model need ~2500 N to change pole angle quickly, so set threshold here.
    if abs(angle) > 0.02:
        if angle >= 0:  
            F_static = 2500  # Static friction coefficient
        else: # input force < 0, so need to reverse sign
            F_static = -2500    
    else: # close to target, make control input small to avoid osc
        F_static = 0
    #print("F_static in dynamics", F_static)    
    
    # Compute derived variables  
    total_mass = m_cart + m_pole  
    sin_theta = np.sin(theta)  
    cos_theta = np.cos(theta)  

    # Equations of motion  
    x_double_dot = (control_input - F_static  + m_pole * length * theta_dot**2 * sin_theta - m_pole * g * sin_theta * cos_theta) / total_mass  
    theta_double_dot = (g * sin_theta - cos_theta * x_double_dot) / (length * (4/3 - (m_pole * cos_theta**2) / total_mass))  

    # Update state  
    next_state = np.array([  
        x + x_dot * dt,  
        x_dot + x_double_dot * dt,  
        theta + theta_dot * dt,  
        theta_dot + theta_double_dot * dt  
    ])  
    return next_state  

# Cost function for MPC  
def cost_function(control_input, initial_state, target_state, N, dt):
    weight_state =  2000  # make state change matters
    weight_input =  0.01
    state = initial_state.copy()  
    total_cost = 0  
    for t in range(N):  
        state = dynamics(state, control_input[t], dt)  
        # Define cost (distance to the target state and control effort)  
        total_cost += weight_state*np.sum(abs(state[2] - target_state[2])) + weight_input*np.sum(abs(control_input[t]))  
    return total_cost  

# MPC Controller  
initial_state = np.array([0.0, 0.0, 0.5, 0.0])  # [x, x_dot, theta, theta_dot]  
target_state = np.array([0.0, 0.0, 0.0, 0.0])  # [target x, target x_dot, target theta, target theta_dot]  
# MPC参数  
N = 3  # 预测步长  
dt = 0.01  
def mpc_control(initial_state, target_state, N=10, dt=0.01):  
    # Initial guess for control inputs  
    control_input_guess = np.zeros(N)  
    #control_input_guess = np.full(N, 2000.0)  # Start with a higher guess
    # Optimization  
    result = minimize(cost_function, control_input_guess, args=(initial_state, target_state, N, dt),  
                      bounds=[(-5000, 5000)] * N)  
    
    # Print full control sequence  
    # print("Full control sequence:", result.x)  
    
    # Print statistics about the control sequence  
    # print("Min control input:", np.min(result.x))  
    # print("Max control input:", np.max(result.x))  
    # print("Mean control input:", np.mean(result.x))  
    
    # Check for small negative values  
    small_neg = result.x[(result.x < 0) & (result.x > -1e-6)]  
    if len(small_neg) > 0:  
        print("Warning: Small negative values detected:", small_neg)  
        
    # Return the first control input  
    return result.x[0]  



# control_signal = mpc_control(initial_state, target_state, N, dt)  
# print(f'Control Signal: {control_signal}')


# 设置仿真参数  
num_steps = 1 
# Simulation loop
try:
    current_state = initial_state.copy() 
    #control_input = 2000
    while True:
        #for control_step in range(20):  
        #state = get_state()  
        
        # 计算控制输入  
        control_input = mpc_control(current_state, target_state, N, dt) # N, dt 
        #control_input = pid_control(state,0.0,0.01)
        #control_input = 100  #200 make cart move and pole balance between drop stand up, by both eye and get state
        print("control input : ", control_input)

        # 应用控制输入  
        p.setJointMotorControl2(cartpole_id, 0, p.TORQUE_CONTROL, force=control_input) 
        
        # Apply control input as an external force to the cart  
        # p.applyExternalForce(objectUniqueId=cartpole_id,   
        # linkIndex=0,  # Base (cart)  
        # forceObj=[control_input, 0, 0],   
        # posObj=p.getBasePositionAndOrientation(cartpole_id)[0],   
        # flags=p.WORLD_FRAME)  
  
        for step in range(num_steps):
            # 前进一步  ,takes time for pole fall, let's wait a little
            p.stepSimulation()
            #time.sleep(0.01)
        
        # cal next state before current_state update 
        # calculate_state = dynamics(current_state, control_input, dt) 
        # print("This is next state by calucation : ", calculate_state)
        #update current state
        current_state = get_state()
        print(f"Step: {step}, Position: {current_state[0]:.2f}, Angle: {current_state[2]:.2f}") 
                           
                
except KeyboardInterrupt:
    # Disconnect from the physics engine
    p.disconnect()


###Problem , pole didn't fall naturally at init angle, only adding force on cart can change
## P control is ok, but need > 250 N to make cart move , is it a problem? 
# it's a problem, make MPC model wrong...
## Tips: 
# MPC control, set bound to 2000 instead of 1000, pole angle become move ( with Pole mass 10kg)
## but when angle move from pos to neg, control force still 2000
## ## user larget timestep to make pole drop faster now, continue to test MPC  -> ok
# use a dynamic updated model to make MPC work, wehn angle < 0.02, F_static = 0 , else, F-static = 2500--> ok
##above method only fits for this cartpole model, when force large enough > 2500 N, pole angle will change
## 200N is balance point that cart  move and pole angle keep same angle ,do not drop
## 2000N make pole quickly change angle
# use angle instead of full states as cost function, control input has negative value, why???
#use angle , the MPC can quickly make pole stand with init pos angle., while -0.5 init pole angle 
#keep at -0.3, why???  --> Friction setting issue, negative force make friction become negative.
#
## even 0.25N make cart move...  --> it is indeed accelerator, with small N, Pole stands quickly, 
# and with less big force added,  cart stop moving.
###since model is not correct, why MPC can choose right control input? --> model may not accurate, but
## increae force can make angle close to 0, this model is correct.
# Amazing MPC !

# print out full control sequence, no negative values  --> ok, use angle state only.

#N > = 3 ok, N = 2 will fail.--> 2 is too short for slow dynamic system, with dt = 0.01, it is 0.02 s behaviour
#increase dt can same to increase N , probably 0.02 s the pole didn't action

#how MPC predict steps, by simulation or calculation?  --> By cacluation, 积分analyze distance, speed, acce

#slidebar mass should set to 0 for fix on GND, it's pybullet rule
##
## Postiion already o --> getbasePositionAndOrientation() get first link postion, which is slidebar, 
## use getLinkState() to get cart position --> ok  -> GPT4 also make mistakes

#by fix above problem, we can check how many force make cart move by get_state() while not naked eye

#MPC robust test, change Pole /cart mass, length etc --> MPC is rubust, change pole mass
#can still make MPC work, it has feedback control and focus on cost func.

###Push into stack
#need 200 N to make it move ...why? --> friction,maybe lower force already move, but eye cannot see.. by get 
#pos it need 200N  --> 力矩，摩擦力，动态力？
#Get system's dynamic func from pybullet data or other tools, through urdf fiel. --> Sympy?
